
-- assembly procedure:
-- 1. convert strings and chars, strip comments and whitespace, check for disallowed chars, split lines, record line numbers, substitute braces (DONE)
-- 2. compute length (WIP)
-- 3. substitute instruction encoding
-- 4. compile to memory dump

local cd = arg[0]:gsub("/", "\\"):gsub("[^\\]+$", "")

local function readfile(fn) local f = io.open(fn, "rb"); assert(f, "could not open file "..fn); local t = f:read("*a"); f:close(); return t; end
local function writefile(fn, d) local f = io.open(fn, "wb"); assert(f, "could not open file "..fn); f:write(d); f:close(); end

function string.trim(s) return s:gsub("^[ \t\r\n]+", ""):gsub("[\t\r\n ]+$", "").."" end
local function loaddefs()
	local defs = "\n"..readfile(cd.."arch.txt"):gsub("\r", "").."\n"
	
	-- apply option lists
	local function applyopts(txt)
		return txt:gsub("\n([^\n%{]*)%{([^%}]+)%}([^\n]*)", function(left, opts, right)
			local l = {}; for opt in opts:gmatch("[^/]+") do
				table.insert(l, (left..opt..right):gsub("%%s1%%", opt):gsub("%%s([0-9]+)%%", function(m) return "%s"..(m-1).."%" end).."")
			end
			return applyopts("\n"..table.concat(l, "\n"))
		end)..""
	end
	defs = applyopts(defs)
	
	local function subpairs(sep, func) defs = defs:gsub("\n([^\n]+)[ \t]*"..sep.."[ \t]*([^\n]+)", func) end
	
	-- apply macros
	local macros = {}; subpairs(":", function(left, right) macros[left:trim()] = right:trim(); return ""; end) -- acquire colon macros
	defs = defs:gsub("[a-zA-Z0-9_]+", macros) -- apply colon macros
	defs = defs:gsub("%$([a-fA-F0-9]+)[ \t]*%|[ \t]*%$([0-9a-fA-F]+)", function(a, b) -- apply OR operations
		a, b = tonumber(a, 16), tonumber(b, 16); assert(bit.band(a, b)==0);
		return "$"..string.format("%02X", bit.bor(a, b))
	end)
	
	local function fixstr(s) return s:trim():gsub(",[ \t]+", ",").."" end
	local instrs = {}; subpairs("=", function(left, right) instrs[fixstr(left)] = fixstr(right) end) -- acquire instructions
	--for k, v in pairs(instrs) do print("["..k.."]", "["..v.."]") end
	return instrs
end

local asminstrs = loaddefs()

local asmerrfn = "?"
local function asmfilename(fn) asmerrfn = fn:match("[^/\\]+$") end
local function asmerror(ln, msg, i) error("Assembly error ("..asmerrfn..", line "..ln.."): "..msg, 1+(i or 0)) end
local function asmassert(ln, v, msg) if not v then asmerror(ln, msg, 1) end; end

local function asmtrimdef(def)
	return def:gsub(";.+$", ""):trim()
end
local function asmmacro(text, macross, macrosf)
	text = text:gsub(".def ([a-zA-Z0-9_]+)%(([a-zA-Z0-9_, \t]+)%)([^\n]+)", function(name, args, def)
		def = asmtrimdef(def)
		if macrosf[name] then asmerror("?", "redefinition of function macro \""..name.."\"") end
		local argc = 0; local argidxs = {}; for arg in args:gmatch("[a-zA-Z0-9_]+") do argc = argc+1; argidxs[arg] = argc; end;
		macrosf[name] = function(argvs)
			if #argvs~=argc then asmerror("?", "wrong number of arguments to function macro \""..name.."\"") end
			return def:gsub("[a-zA-Z0-9_]+", function(a) return argvs[argidxs[a]] end):gsub("%%", "")
		end
		return ""
	end)
	text = text:gsub(".def ([a-zA-Z0-9_]+)([^\n]+)", function(name, def)
		if macross[name] then asmerror("?", "redefinition of string macro \""..name.."\"") end
		macross[name] = asmtrimdef(def)
		return ""
	end)
	
	text = text:gsub("([a-zA-Z0-9_]+)", macross)
	text = text:gsub("([a-zA-Z0-9_]+)%(([^%)]+)%)", function(name, args)
		if macrosf[name] then
			local argvs = {}; for arg in args:gmatch("[^,]+") do table.insert(argvs, arg:trim().."") end;
			return macrosf[name](argvs)
		end
	end)
	
	return text
end

local escapes = { ["\""] = "\"", ["\'"] = "\'", ["\\"] = "\\", ["0"] = "\0", ["t"] = "\t", ["r"] = "\r", ["n"] = "\n", }
local function flipt(s) local t = {}; for i = 1, #s do local c = s:sub(i, i); t[c] = true; end; return t; end
local normalchars = flipt("abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789_() .$[],#")
local function asmclean(text, fid)
	local str, esc, com, spc, stk, brc, t2t, nel, par, lin = false, false, false, false, {}, (fid or 1)*10000, {}, false, 0, 1
	local function newline() if nel then table.insert(t2t, "\n"); nel = false; end end
	local function insert(t) if (not nel) then table.insert(t2t, lin.."\t") end; table.insert(t2t, t); nel = true; end
	for i = 1, #text do local c = text:sub(i, i)
		if c=="\n" then lin = lin+1 end
		if str then
			if esc then insert("$"..string.format("%02X", string.byte(escapes[c] or asmerror(lin, "invalid escape "..c)))); newline(); esc = false;
			else
				if c=="\\" then esc = true
				elseif c=="\"" or c=="\'" then str = false
				else insert("$"..string.format("%02X", string.byte(c))); newline();
				end
			end
		elseif com then
			if c=="\n" then com = false; newline(); end
		else
			if c=="\"" or c=="\'" then str = true
			elseif c==";" then com = true
			elseif c=="\r" or ((c==" " or c=="\t") and (spc or (not nel))) then
			elseif c==" " or c=="\t" then table.insert(t2t, " ");
			elseif c=="{" then if nel then insert("_BRACE_"..brc.."_CLOSE") end; newline(); insert("_BRACE_"..brc.."_OPEN:"); newline(); table.insert(stk, brc); brc = brc+1;
			elseif c=="}" then if nel then insert("_BRACE_"..(stk[#stk]).."_OPEN") end; newline(); insert("_BRACE_"..(stk[#stk]).."_CLOSE:"); newline(); table.remove(stk, #stk);
			elseif c=="\n" or (c=="," and par==0) then newline();
			elseif c==":" then insert(":"); newline();
			elseif c=="(" then par = par+1; insert(c);
			elseif c==")" then par = par-1; insert(c);
			elseif normalchars[c] then insert(c);
			else asmerror(lin, "invalid character "..c)
			end
			spc = (c==" " or c=="\t")
		end
	end
	return table.concat(t2t)
end

local function asmmatch(ln, line, labelnames)
	line = line:gsub("%$[a-fA-F0-9][a-fA-F0-9][a-fA-F0-9][a-fA-F0-9]", "%%imm16%%")
	line = line:gsub("%$[a-fA-F0-9][a-fA-F0-9]"                      , "%%imm8%%" )
	return asminstrs[line] or
		asminstrs[line:gsub("[a-zA-Z0-9_]+", function(n) if labelnames[n] then return "%imm16%"   end end)] or
		asmerror(ln, "No instruction matching line \""..line.."\"")
end
local function asmnumber(n)
	return tonumber(n:sub(2, #n), 16) or asmerror("?", "invalid number \""..n.."\"")
end
local function asmalign(addr, aln)
	if addr%aln==0 then return addr
	else return addr+(aln-(addr%aln))
	end
end

-- Assemble: Convert cleaned assembly text into 
local function asmcompile(text, ram)
	
	local labelnames = {}; for name in text:gmatch("\t([a-zA-Z0-9_]+):\n") do labelnames[name] = true end;
	
	local addr = 0
	local lines = {}
	local labels = {}
	for ln, line in text:gmatch("([0-9]+)\t([^\n]+)") do
		line = line:gsub(" +$", "")
		ln = tonumber(ln)
		if line:find(":$") then -- label
			local name = line:sub(1, #line-1)
			if labels[name] then asmerror(ln, "Label redefined: \""..name.."\"") end
			labels[name] = "$"..string.format("%04X", addr)
		elseif line:find("^%$") then -- data
			local len = (#line-1)/2
			table.insert(lines, { ln = ln, addr = addr, ins = line..",", text = line, len = len, mul = 1, })
			addr = addr+len
			
		elseif line:find("^%.org") then -- org
			addr = asmnumber(line:sub(6, #line))
		elseif line:find("^%.aln") then -- align
			addr = asmalign(addr, asmnumber(line:sub(6, #line)))
		elseif line:find("^%.fil") then -- filename
			asmfilename(line:sub(6, #line))
		else -- instruction
			local mul = 1; line = line:gsub("%[([^%]]+)%]", function(i) mul = mul*asmnumber(i:trim()); return ""; end)
			line = line:gsub(" +$", "")
			local ins = asmmatch(ln, line, labelnames)
			local len = 0; for _ in ins:gmatch(",") do len = len+1 end
			table.insert(lines, { ln = ln, addr = addr, ins = ins, text = line, len = len, mul = mul, })
			addr = addr+len*mul
		end
	end
	
	local function setram(ln, a, v)
		if ram[a] then asmerror(ln, "Overwriting memory at %04X", a) end
		assert(v==true or (type(v)=="number" and v>=0 and v<=255), "ram value = "..tostring(v))
		ram[a] = v
	end
	for _, line in ipairs(lines) do
		local text = line.text:gsub("[a-zA-Z0-9_]+", labels)
		local imms = {}; text:gsub("$[0-9a-fA-F]+", function(n) table.insert(imms, asmnumber(n)) end)
		local ins = line.ins:gsub("%%i([0-9])([a-z]?)%%", function(idx, lh)
			local imm = imms[tonumber(idx)]
			if lh=="l" then imm = imm%256
			elseif lh=="h" then imm = math.floor(imm/256)
			elseif lh=="r" then imm = imm-(line.addr+line.len); if imm<-128 or imm>127 then asmerrror(line.ln, "relative jump too far ("..imm..")") end; imm = imm%256;
			end
			if imm<0 or imm>255 then asmerror(line.ln, "immediate too large ("..imm..")") end
			return "$"..string.format("%02X", imm)
		end)
		
		local len = 0
		for n in ins:gmatch("[^,]+") do
			local a = line.addr+len; len = len+1;
			setram(line.ln, a, (n=="?" and true) or asmnumber(n))
		end
		assert(len==line.len)
		for i = 1, line.mul-1 do
			for j = 0, len-1 do
				setram(line.ln, line.addr+(i*len)+j, ram[line.addr+j])
			end
		end
	end
end

local function assemble(fn1, ram)
	local inctext = {}
	local incname = {}
	local incnames = {}
	local function getincs(ifn)
		local text = readfile(ifn)
		return text:gsub(".inc \"([^\"]+)\"", function(fn)
			if not incnames[fn] then
				incnames[fn] = true
				local idx = #inctext+1
				local it = getincs(fn)
				inctext[idx] = it
				incname[idx] = fn
				return ".inc "..idx
			else
				asmfilename(ifn)
				asmerror("?", "File included multiple times: \""..fn.."\"")
				return ""
			end
		end)..""
	end
	incnames[fn1] = true
	local maintext = getincs(fn1)
	
	-- apply macros
	local macross, macrosf = {}, {}
	asmfilename(fn1); maintext = asmmacro(maintext, macross, macrosf)
	for i, text in ipairs(inctext) do asmfilename(incname[i]); inctext[i] = asmmacro(text, macross, macrosf) end;
	asmfilename(fn1); maintext = asmclean(maintext, 1)
	for i, text in ipairs(inctext) do asmfilename(incname[i]); inctext[i] = "0\t.fil "..incname[i].."\n"..asmclean(text, i+1).."0\t.fil "..fn1.."\n" end -- clean text
	
	maintext = maintext:gsub("[0-9]+\t.inc ([0-9]+)\n", function(idx)
		idx = tonumber(idx)
		return inctext[idx]
	end)
	
	--print(maintext)
	
	asmcompile(maintext, ram)
end

----

local function printablechar(c)
	if c>=0x20 and c<=0x7E then return string.char(c)
	else return "." end
end
local function memdump(ram)
	local printblanks = false
	
	local tt = {}
	local function ins(s) table.insert(tt, s) end
	ins("Memory dump:")
	ins("Addr  |   0  1  2  3  4  5  6  7   8  9  A  B  C  D  E  F  |  ")
	local lastlineempty = true
	for laddr = 0x0000, 0xFFF0, 0x10 do
		local empty = true; for addr = laddr, laddr+0x0F do if (ram[addr]==true and printblanks) or (type(ram[addr])=="number") then empty = false end end;
		--if (not empty) or laddr==0x0000 or laddr==0xFFF0 then
		if not empty then
			if lastlineempty then ins("      |") end
			local line = {}
			local linec = {}
			for addr = laddr, laddr+0x0F do
				if ram[addr]==true and printblanks then table.insert(line, "??"); table.insert(linec, "?")
				elseif type(ram[addr])=="number" then table.insert(line, string.format("%02X", ram[addr])); table.insert(linec, printablechar(ram[addr]));
				else table.insert(line, "--"); table.insert(linec, "-")
				end
			end
			table.insert(line, 9, "")
			ins(string.format("%04X", laddr).."  |  "..table.concat(line, " ").."  |  "..table.concat(linec))
		end
		
		lastlineempty = empty
	end
	ins("")
	return table.concat(tt, "\n")
end

local function memoutput(ram)
	local data = {}
	local addr = 0
	local function writebyte(b, pos) if not pos then addr = addr+1 end; pos = pos or addr-1; assert(b>=0 and b<=0xFF and b%1==0); data[pos] = b; return pos; end
	local function writeword(w, pos) if not pos then addr = addr+2 end; pos = pos or addr-2; assert(w>=0 and w<=0xFFFF and w%1==0); data[pos] = math.floor(w/256); data[pos+1] = w%256; return pos; end
	
	local lenpos = 0
	local len = 0
	for a = 0x0000, 0xFFFF do
		if type(ram[a])=="number" then
			if len==0 then writeword(a); lenpos = writeword(0); end
			len = len+1
			writebyte(ram[a])
		else
			if len>0 then writeword(len, lenpos) end
			len = 0
		end
	end
	if len>0 then writeword(len, lenpos) end
	
	local datas = {}; for i = 0, addr-1 do datas[i+1] = string.char(data[i]) end
	return table.concat(datas)
end

----

local helptext =
[[8700 Assembler
	See doc.txt for assembly syntax information
	Usage: asm.lua -i input_file [-i other_input_files...] [-o output_files] [-d dump_files]
	
]]

local input_files = {}
local dump_targets = {}
local output_files = {}

local i = 1; while arg[i] do
	local a, a2 = arg[i], arg[i+1]
	i = i+1
	if a=="-d" or a=="--dump" then
		assert(a2, "--dump flag requires a target")
		table.insert(dump_targets, a2)
		i = i+1
	elseif a=="-i" or a=="--input" then
		assert(a2, "--input flag requires a target")
		table.insert(input_files, a2)
		i = i+1
	elseif a=="-o" or a=="--output" then
		assert(a2, "--output flag requires a target")
		table.insert(output_files, a2)
		i = i+1
	elseif a=="-h" or a=="--help" then
		print(helptext)
	else
		error("invalid flag "..a)
	end
end
if #arg==0 then print(helptext) end

local ram = {}

if #input_files>0 then
	for _, fn in ipairs(input_files) do
		if fn=="stdin" then
			local text = io.read("*a")
			assemble(text, ram)
		else
			assemble(fn, ram)
		end
	end
else
	error("No input files")
end

if #dump_targets>0 then
	local md = memdump(ram)
	for _, fn in ipairs(dump_targets) do
		if fn=="stdout" then
			print(md)
		else
			writefile(fn, md)
		end
	end
end

if #output_files>0 then
	local out = memoutput(ram)
	for _, fn in ipairs(output_files) do
		if fn=="stdout" then
			io.write(out)
		elseif fn=="-stdout" then
			local t = {}; for i = 1, #out do table.insert(t, string.format("%02x", string.byte(out:sub(i, i)))) end;
			print(table.concat(t))
		else
			writefile(fn, out)
		end
	end
end
